import os
from dashscope import MultiModalConversation
import csv
image_dir = "/Users/muzhiyi/Desktop/perscription_senior/perscription-basic"
image_files = [f for f in os.listdir(image_dir) if f.endswith(('.png', '.jpg',  '.jpeg'))]
import os
import csv

# CSV 文件路径
csv_file_path = '/Users/muzhiyi/Desktop/perscription_senior/output_basic/base_qwen2-vl-7b-instruct_output.csv'

if not os.path.exists(csv_file_path):
    with open(csv_file_path, 'w', newline='', encoding='utf-8') as csvfile:
        fieldnames = ['filename', 'description', 'analysis']
        writer = csv.DictWriter(csvfile, fieldnames=fieldnames)
        writer.writeheader()
else:
    pass

import base64

def image_to_data_url(file_path):
    with open(file_path, "rb") as image_file:
        encoded_str = base64.b64encode(image_file.read()).decode("utf-8")
    mime_type = "image/jpeg" if file_path.lower().endswith((".jpg", ".jpeg")) else "image/png"
    return f"data:{mime_type};base64,{encoded_str}"


# import pandas as pd
#
# # 尝试读取 CSV，若不存在则初始化 DataFrame
# try:
#     df = pd.read_csv(csv_file_path)
# except FileNotFoundError:
#     df = pd.DataFrame(columns=['filename', 'description', 'analysis'])

# # 添加新数据
# new_data = pd.DataFrame([{'filename': 'test.png', 'description': '...', 'analysis': '...'}])
# df = pd.concat([df, new_data], ignore_index=True)

# 保存回 CSV
# df.to_csv(csv_file_path, index=False)

def analyze_prescription(image_file):
    from openai import OpenAI
    import os

    client = OpenAI(
        # 若没有配置环境变量，请用百炼API Key将下行替换为：api_key="sk-xxx"
        api_key="sk-26410c86527540879526c1c95454f345",
        base_url="https://dashscope.aliyuncs.com/compatible-mode/v1"
    )
    messages = [
        {
            "role": "system",
            "content": [{"type": "text", "text": "You are a helpful assistant."}]},
        {
            "role": "user",
            "content": [
                {
                    "type": "image_url",
                    "image_url": {
                        "url": image_to_data_url(image_file)
                    },
                },
                {"type": "text", "text": "识别图中处方内容， 要求给出{{姓名：}，{性别：},{年龄：},{科室：},{时间：},{诊断：},{药物1名称： ,用法：},{药物2名称: ， 用法： }（其他药物写法以此类推）...,{医师：},{药师：},{审核： },{调配： },{核对： },{发药： },{金额： }}等信息"},
            ],
        }
    ]
    completion1 = client.chat.completions.create(
        model="qwen2-vl-7b-instruct",
        messages=messages,
    )
    print(f"第一轮输出：{completion1.choices[0].message.content}")
    output1 = completion1.choices[0].message.content
    assistant_message = completion1.choices[0].message
    messages.append(assistant_message.model_dump())
    messages.append({
        "role": "user",
        "content": [
            {
                "type": "text",
                "text": "关于给出的处方内容，结合现有的医学知识，识别药物搭配是否合理，是否存在剂量过多的情况。 如果是安全的，则回答safe， 如果不合理， 则回答unsafe，并给出相应原因。回答格式为：先结论：Safe/Unsafe。 后原因：Reason:....."
            }
        ]
    })
    completion2 = client.chat.completions.create(
        model="qwen2-vl-7b-instruct",
        messages=messages,
    )
    print(f"第二轮输出：{completion2.choices[0].message.content}")
    output2 = completion2.choices[0].message.content
    return   output1,output2

# 后续的写入操作可以使用 'a'（追加）模式或 'w'（覆盖）模式
with open(csv_file_path, 'a', newline='', encoding='utf-8') as csvfile:
    fieldnames = ['filename', 'description', 'analysis']
    writer = csv.DictWriter(csvfile, fieldnames=fieldnames)

    for image_file in image_files:
        image_path = os.path.join(image_dir, image_file)
        output1, output2 = analyze_prescription(image_path)
        writer.writerow({'filename': image_file, 'description': output1, 'analysis': output2})


